import torch
from torch import nn
from darknet53 import DarkNet53
from config.module import ConvolutionSet, CBL, UpSample


class YoLov3(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = DarkNet53()
        self.conv1 = nn.Sequential(
            ConvolutionSet(1024, 512)
        )

        self.detect_13 = nn.Sequential(
            CBL(512, 1024, 3, 1),
            # 4分类 * 锚框： (1 + 4 + 4) * 3
            CBL(1024, 27, 1, 1)
        )

        self.neck_hidden1 = nn.Sequential(
            CBL(512, 256, 1, 1),
            UpSample()
        )

        self.conv2 = nn.Sequential(
            ConvolutionSet(256 + 512, 256)
        )

        self.detect_26 = nn.Sequential(
            CBL(256, 512, 3, 1),
            CBL(512, 27, 1, 1)
        )

        self.neck_hidden2 = nn.Sequential(
            CBL(256, 128, 1, 1),
            UpSample()
        )

        self.conv3 = nn.Sequential(
            ConvolutionSet(128 + 256, 128)
        )

        self.detect_52 = nn.Sequential(
            CBL(128, 256, 3, 1),
            CBL(256, 27, 1, 1)
        )

    def forward(self, x):
        backbone_unit52_out, backbone_unit26_out, backbone_unit13_out = self.backbone(x)
        conv_13_out = self.conv1(backbone_unit13_out)

        neck_13_26_out = self.neck_hidden1(conv_13_out)
        route26_out = torch.cat((neck_13_26_out, backbone_unit26_out), dim=1)

        conv_26_out = self.conv2(route26_out)

        neck_26_52_out = self.neck_hidden2(conv_26_out)
        route52_out = torch.cat((neck_26_52_out, backbone_unit52_out), dim=1)

        conv_52_out = self.conv3(route52_out)

        detect_13_out = self.detect_13(conv_13_out)
        detect_26_out = self.detect_26(conv_26_out)
        detect_52_out = self.detect_52(conv_52_out)
        return detect_13_out, detect_26_out, detect_52_out


if __name__ == '__main__':
    data = torch.randn(1, 3, 416, 416)
    yolov3 = YoLov3()
    outs = yolov3(data)
    for out in outs:
        print(out.shape)
    # torch.Size([1, 512, 13, 13])
    # torch.Size([1, 256, 26, 26])
    # torch.Size([1, 128, 52, 52])
    # P0 P1 P2:  N  27  H  W
    # 27: 类别(1 + 4 + 4) * 锚框3
    # torch.Size([1, 27, 13, 13])
    # torch.Size([1, 27, 26, 26])
    # torch.Size([1, 27, 52, 52])
    pass
